/* 
 * Quechua - the lightweight data mining framework
 *
 * Copyright (C) 2012 Marek Denis <quechua@octogan.net>
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program.  If not, see <http://www.gnu.org/licenses/>.
 */

#ifndef IPPROC_PROCESSOR_H
#define IPPROC_PROCESSOR_H

#include <sys/types.h>
#include <sys/socket.h>
#include <netdb.h>
#include <arpa/inet.h>

#include "../../../include/interface-processor.h"
#include "../../data/dionaeaconnections.h"

#include "../../data/itemsets.h"
#include "../../data/indextransactions.h"

#include "attributes.h"
#include "iptrie.h"

using std::nothrow;

class IpProcessor : public Processor {
 public:
    IpProcessor();
    IpProcessor(u_int64_t dst, u_int64_t src); 
    virtual ~IpProcessor();
    datapack_ptr preprocess(datapack_ptr data);
    datapack_ptr postprocess(datapack_ptr);
    virtual bool reset();

 protected:
    /* internal methods*/
    index_t misc_to_index(const char* misc,const u_int16_t val);
    index_t ip_to_index(const char* ip,INDEX_TYPE dir);
    const char* index_to_ip(const index_t& idx);
    in_addr*  encode_ip4(const char* ip, in_addr* buffer);
    in6_addr* encode_ip6(const char* ip, in6_addr* buffer);
    const char* decode_ip4(in6_addr* buffer,char* ip);
    const char* decode_ip6(in6_addr* buffer,char* ip);
    const char* decode_ip4_fc(const u_int8_t* buffer,char* ip);
    const char* decode_ip6_fc(const u_int8_t* buffer,char* ip);

    index_t get_index_from_tree(IpTrie* tree ,u_int8_t* label,u_int8_t len,const char* ip, INDEX_TYPE dir,IpNode** ext_node);

    u_int8_t ip_family(const char* ip);
    index_t allocate_src_idx();
    index_t allocate_dst_idx();
    /*******************/
    void init();
    void debug();
    void clear();

    Attributes* attributes;
    IpTrie* ipv4;
    IpTrie* ipv6;
    IpTrie* rev_tree;

    index_t src_idx;
    index_t dst_idx;

    index_t misc_maxval;
    index_t dst_maxval;
    index_t src_maxval;
};

IpProcessor::IpProcessor() : Processor() {
    dst_idx = 131078; // 0x20006
    src_idx = 9223372036854710267; // 0x7ffffffffffefffc

    misc_maxval = dst_idx - 1;
    dst_maxval  = src_idx - 1;
    src_maxval  = 0xffffffffffffffff; 

    init();
};

IpProcessor::IpProcessor(u_int64_t dst, u_int64_t src) : Processor() {
    dst_idx = dst;
    src_idx = src;

    misc_maxval = dst_idx - 1;
    dst_maxval  = src_idx - 1;
    src_maxval  = 0xffffffffffffffff; 

    init();
}

IpProcessor::~IpProcessor() {
    clear();
};

void IpProcessor::clear() {
    delete attributes;
    delete ipv4;
    delete ipv6;
    delete rev_tree;

};

void IpProcessor::init() {
    attributes = new(nothrow) Attributes();
    if(!attributes) {
        LOG(ERROR) << "IpProcessor: Error while initializing Attributes object";
        return;
    }

    attributes->set_attribute(SINGLE_ATTRIBUTE,"tcp",1);
    attributes->set_attribute(SINGLE_ATTRIBUTE,"udp",1);
    attributes->set_attribute(RANGE_ATTRIBUTE,"local_port",65535);
    attributes->set_attribute(RANGE_ATTRIBUTE,"remote_port",65535);


    ipv4 = new(nothrow) IpTrie(4);

    if(!ipv4) {
        delete attributes;
        LOG(ERROR) << "IpProcessor: Error while initializing ipv4 tree";
        return;
    }

    ipv6 = new(nothrow) IpTrie(16);
    
    if(!ipv6) {
        delete attributes;
        delete ipv4;
        LOG(ERROR) << "IpProcessor: Error while initializing ipv6 tree";
        return;
    }

    rev_tree = new(nothrow) IpTrie(sizeof(index_t));

    if(!rev_tree) {
        delete attributes;
        delete ipv4;
        delete ipv6;
        LOG(ERROR) << "IpProcessor: Error while initializing rev_tree";
        return;
    }
};


/*
 * Indexes non-ip address to index_t value (ports,protocols, flags etc)
 */
index_t IpProcessor::misc_to_index(const char* misc, const u_int16_t val) {
    return attributes->get_attribute_index(misc,val);
};

/*
 * Find ip address family
 * encode to binary data
 * create/find ipIndex node and save it if necessary.
 * return index_t value (found/generated)
 */ 

index_t IpProcessor::ip_to_index(const char* ip,INDEX_TYPE dir) {
    IpTrie* tree;
    IpNode** node = new IpNode*;
    index_t idx;
    //LOG(CRITIC) << "Indexing ip: " << ip;
    switch(ip_family(ip)) {
        case AF_INET:
            tree = ipv4;
            // think about getting rid of memcpy
            in_addr a4;
            encode_ip4(ip,&a4);
            idx =  get_index_from_tree(tree,(u_int8_t*)&(a4.s_addr),4,ip,dir,node); 
            rev_tree->setNode((const u_int8_t*)&idx,sizeof(idx),(*node)->value);
            break;
        case AF_INET6:
            tree = ipv6;
            in6_addr a6;
            encode_ip6(ip,&a6);
            idx = get_index_from_tree(tree,(u_int8_t*)&(a6.s6_addr),16,ip,dir,node);
            rev_tree->setNode((const u_int8_t*)&idx,sizeof(idx),(*node)->value);
            break;
        default:
            LOG(ERROR) << "Unknown family for address " << ip;
            return 0;
    };
    delete node;
    return idx;
}

u_int8_t IpProcessor::ip_family(const char* ip) {
    struct addrinfo hint, *info = NULL;
    memset(&hint,0,sizeof(hint));
    hint.ai_family = AF_UNSPEC;
    hint.ai_flags = AI_NUMERICHOST; // turn off DNS lookups
    if(getaddrinfo(ip,0,&hint,&info)) { // returns 0 on success
        LOG(ERROR) << "IpProcessor: Problems with getaddrinfo(3), ip: " << ip;
        return AF_UNSPEC;
    }
    u_int8_t f = info->ai_family;
    freeaddrinfo(info);
    return f;
};

index_t IpProcessor::allocate_src_idx() {
    return src_idx++;
};
index_t IpProcessor::allocate_dst_idx() {
    return dst_idx++;
};

in_addr* IpProcessor::encode_ip4(const char* ip, in_addr* addr) {
    int del = 0;
    if(!addr) {
        if(!(addr = new(std::nothrow) in_addr)) {
            LOG(ERROR) << "IpProcessor: Error while initializing buffer for ip4 address";
            return addr;
        }
        del = 1;
    }
    if(!(inet_pton(AF_INET,ip,addr))) {
        if(del) { 
            delete addr;
            addr = NULL;
        }
    }
    return addr;
};


in6_addr* IpProcessor::encode_ip6(const char* ip, in6_addr* addr) {
    int del = 0;
    if(!addr) {
        if(!(addr = new(std::nothrow) in6_addr)) {
            LOG(ERROR) << "IpProcessor: Error while initializing buffer for ip6 address";
            return addr;
        }
        del = 1;
    }
    if(!(inet_pton(AF_INET6,ip,addr))) {
        if(del) { 
            delete addr;
            addr = NULL;
        }
    }
    return addr;
};

const char* IpProcessor::decode_ip4(in6_addr* buffer,char* ip) {
    int del = 0;
    if(!ip) {
        if(!(ip = new (std::nothrow) char[INET_ADDRSTRLEN] )) {
            LOG(ERROR) << "IpProcessor: Error while initializing buffer for ip4 address";
            return ip;
        }
        del = 1;
    } 
    
    // was inet_ntop(AF_INET,(void*)buffer->s_addr,ip,4)
    u_int8_t* data = buffer->s6_addr;
    if(!(sprintf(ip,"%u.%u.%u.%u",data[0], data[1], data[2], data[3]))) {
       if(del)
           delete ip;
       LOG(ERROR) << "IpProcessor: Error while decoding IPv4 address";
       return NULL;
    } 
    return ip;
};

const char* IpProcessor::decode_ip6(in6_addr* buffer,char* ip) {
    int del = 0;
    if(!ip) {
        if(!(ip = new (std::nothrow) char[INET6_ADDRSTRLEN] )) {
            LOG(ERROR) << "IpProcessor: Error while initializing buffer for ip4 address";
            return ip;
        }
        del = 1;
    }

    //inet_ntop(AF_INET6,(void*)buffer->s6_addr,ip,16)
    u_int8_t* data = buffer->s6_addr;
    int chars = sprintf(ip,"%x:%x:%x:%x:%x:%x:%x:%x",
                       (data[0] << 8) | data[1], (data[2] << 8) | data[3],
                       (data[4] << 8) | data[5], (data[6] << 8) | data[7],
                       (data[8] << 8) | data[9], (data[10] << 8) | data[11],
                       (data[12] << 8) | data[13], (data[14] << 8) | data[15]);
    if(!chars) {
       if(del)
           delete ip;
       LOG(ERROR) << "IpProcessor: Error while decoding IPv4 address";
       return NULL;
    } 
    return ip;

};

const char* IpProcessor::decode_ip4_fc(const u_int8_t* buffer,char* ip) {
    int del = 0;
    if(!ip) {
        if(!(ip = new (std::nothrow) char[INET_ADDRSTRLEN] )) {
            LOG(ERROR) << "IpProcessor: Error while initializing buffer for ip4 address";
            return ip;
        }
        del = 1;
    }

//    in_addr addr;
//    memcpy((void*)addr.s_addr,(void*)buffer,4);
//    was(inet_ntop(AF_INET6,buffer,ip,16))
    if(!(sprintf(ip,"%u.%u.%u.%u",buffer[0], buffer[1], buffer[2], buffer[3]))) {
       if(del)
           delete ip;
       LOG(ERROR) << "IpProcessor: Error while decoding IPv4 address";
       return NULL;
    } 
    return ip;


};

const char* IpProcessor::decode_ip6_fc(const u_int8_t* buffer,char* ip) {
    int del = 0;
    if(!ip) {
       if(!(ip = new (std::nothrow) char[INET6_ADDRSTRLEN] )) {
           LOG(ERROR) << "IpProcessor: Error while initializing buffer for ip4 address";
           return ip;
       }
       del = 1;
    }

//    in6_addr addr;
//    memcpy((void*)addr.s6_addr,(void*)buffer,16);
//    if(!(inet_ntop(AF_INET6,buffer,ip,16))) {
    int chars = sprintf(ip,"%x:%x:%x:%x:%x:%x:%x:%x",
                       (buffer[0] << 8) | buffer[1], (buffer[2] << 8) | buffer[3],
                       (buffer[4] << 8) | buffer[5], (buffer[6] << 8) | buffer[7],
                       (buffer[8] << 8) | buffer[9], (buffer[10] << 8) | buffer[11],
                       (buffer[12] << 8) | buffer[13], (buffer[14] << 8) | buffer[15]);

    if(!chars) {
       if(del)
           delete ip;
       LOG(ERROR) << "IpProcessor: Error while decoding IPv6 address";
       return NULL;
    } 
    return ip;
};

index_t IpProcessor::get_index_from_tree(IpTrie* tree ,u_int8_t* label,u_int8_t len,const char* ip, INDEX_TYPE dir,IpNode** extnode) {

    IpNode* node;

    if(!(node = tree->getNode(label,len))) {// didn't find any node
        ipIndex_ptr idx_ptr(new ipIndex);

        if(!(node = tree->setNode(label,len,idx_ptr,false))) {
            LOG(ERROR) << "IpProcessor: Cannot allocate new IpNode";
            return 0;
        }
        *extnode = node;
        if(!ip) { // argument was set to NULL, need to handle memory in ipIndex.
            if(len==4)
                ip = decode_ip4_fc(label,NULL);
            else if(len==16)
                ip = decode_ip6_fc(label,NULL);
            else 
                ip = "255.255.255.255";

            idx_ptr->handle_mem(true);
        }
        //LOG(CRITIC) << "get_index_from tree(): we got ip: " << ip;
        idx_ptr->ip = ip;

        idx_ptr->ver = len==4?AF_INET:AF_INET6; // idx->ver
        if(len!=4 && len !=16) {
            LOG(ERROR) << "get_index_from_tree: len is neither 4 nor 16, len:" << len;
        }
        idx_ptr->type= dir;                     // idx->type
        memcpy((void*)idx_ptr->addr->s6_addr,(void*)label,len); //idx->addr
        if(idx_ptr->type == SRC) {
            idx_ptr->src_idx = allocate_src_idx();                   //idx->src_idx
            return idx_ptr->src_idx;
        }
        if(idx_ptr->type == DST) {
            idx_ptr->dst_idx = allocate_dst_idx();                   //idx->dst_idx
            return idx_ptr->dst_idx;
        }
    } 
    // we got here - the node was found
    
    ipIndex_ptr idx_ptr = node->value;
    *extnode = node;

    if(idx_ptr->type == BOTH) {
        if(dir==SRC) return idx_ptr->src_idx;
        if(dir==DST) return idx_ptr->dst_idx;
        return 0; //error
    }

    if(idx_ptr->type == SRC) {
        if(dir==SRC)
            return idx_ptr->src_idx;
        // dir == DST
        idx_ptr->dst_idx = allocate_dst_idx();
        idx_ptr->type = BOTH;
        return idx_ptr->dst_idx;
    }
    
    if(idx_ptr->type == DST) {
        if(dir==DST)
            return idx_ptr->dst_idx;
        // dir == SRC
        idx_ptr->src_idx = allocate_src_idx();
        idx_ptr->type = BOTH;
        return idx_ptr->src_idx;
    }
}

const char* IpProcessor::index_to_ip(const index_t& idx) {
    IpNode* node = NULL;
    const char* ip;
    if(!(node = rev_tree->getNode((u_int8_t*)&idx,(u_int8_t)sizeof(index_t)))) {
        LOG(ERROR) << "IpProcessor:index_to_ip():: didn't find node for idx: " << idx;
        return NULL;
    }
    ipIndex_ptr idx_node = node->value;
    return idx_node->ip;
};


datapack_ptr IpProcessor::preprocess(datapack_ptr data) {
    DionaeaConnections* connections = dynamic_cast<DionaeaConnections*>(data.get());
    IndexTransactions* transactions = new IndexTransactions();

    list<DionaeaConnections::unindexedTransaction*>::iterator it;
    DionaeaConnections::unindexedTransaction* txn;
    
    itemset* itxn = NULL;

    it = connections->rows.begin();
    int i=1;
    for(;it!=connections->rows.end();++it,++i) {
        txn = *it;
#undef DEBUG_MSG
#ifdef DEBUG_MSG
        //proto
        index_t proto_idx = attributes->get_attribute_index(txn->protocol,1);
        // ports
        index_t lport_idx = attributes->get_attribute_index("local_port",txn->dst_port);
        index_t rport_idx = attributes->get_attribute_index("remote_port",txn->src_port);
        // hosts
        index_t lhost_idx = ip_to_index(txn->dst_host,DST);
        index_t rhost_idx = ip_to_index(txn->src_host,SRC);

        printf("Transaction %d : %s %s:%d->%s:%d ==> %lu %lu:%lu->%lu:%lu\n",i,
                txn->protocol,txn->src_host,txn->src_port,
                txn->dst_host,txn->dst_port,
                proto_idx,
                rhost_idx,rport_idx,
                lhost_idx,lport_idx
                );
        fflush(stdout);
#endif

        itxn = new itemset(5);
        itxn->val(0) = attributes->get_attribute_index(txn->protocol,1);
        itxn->val(1) = attributes->get_attribute_index("local_port",txn->dst_port);
        itxn->val(2) = attributes->get_attribute_index("remote_port",txn->src_port);
        itxn->val(3) = ip_to_index(txn->dst_host,DST);
        itxn->val(4) = ip_to_index(txn->src_host,SRC);
        transactions->append(itxn);

    }
    return datapack_ptr(transactions);
}

datapack_ptr IpProcessor::postprocess(datapack_ptr data) {
    IndexTransactions* itemsets = dynamic_cast<IndexTransactions*>(data.get());
    DionaeaConnections* connections = new DionaeaConnections();

    DionaeaConnections::unindexedTransaction* con;
    
    itemset* item;

    for(int j=0;j<itemsets->getsize();++j) {
        item = itemsets->getitem(j);
        con = new DionaeaConnections::unindexedTransaction();

        con->count = item->getcount();
        con->generator = item->getgenerator();
        
        for(int i=0;i<item->getsize();++i) {
            attribute_val attr_val;
            if(item->val(i) <= misc_maxval) {
                attributes->get_attribute_value(item->val(i),&attr_val);
                if(!string_t(attr_val.desc).compare("local_port")) {
                    con->dst_port = attr_val.value;
                }
                else if(!string_t(attr_val.desc).compare("remote_port")) {
                    con->src_port = attr_val.value;
                }
                else {
                    con->protocol = attr_val.desc;
                }
            } 
            else if(item->val(i) <= dst_maxval) {
                con->dst_host = index_to_ip(item->val(i));
            } 
            else {
                con->src_host = index_to_ip(item->val(i));
            }
        }
        connections->rows.push_back(con);
    }
    itemsets->clear();
    
    return datapack_ptr(connections);
}

bool IpProcessor::reset() {
    clear();
    init();
    return true;
}

extern "C" Processor* create_processor() {
    return new IpProcessor();
};

#endif
